/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * openUBMC is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

#include <gtest/gtest.h>

#include <mc/runtime/thread_list.h>
#include <mc/sync/shared_mutex.h>

#include <future>
#include <shared_mutex>

#include "../runtime/test_future_helpers.h"

using namespace mc::sync;

template <typename Mutex>
class UpgradeMutexTest : public ::testing::Test {
public:
    using MutexType = Mutex;
};

using MutexTypes = ::testing::Types<shared_mutex, reader_priority_shared_mutex>;
TYPED_TEST_SUITE(UpgradeMutexTest, MutexTypes);

// 基本升级锁功能测试
TYPED_TEST(UpgradeMutexTest, BasicUpgradeLock) {
    TypeParam m;

    // 获取升级锁
    m.lock_upgrade();
    m.unlock_upgrade();
}

TYPED_TEST(UpgradeMutexTest, TryUpgradeLock) {
    TypeParam m;

    ASSERT_TRUE(m.try_lock_upgrade());
    m.unlock_upgrade();

    // 升级锁与升级锁互斥
    m.lock_upgrade();
    ASSERT_FALSE(m.try_lock_upgrade());
    m.unlock_upgrade();
}

TYPED_TEST(UpgradeMutexTest, UpgradeLockWithTimeout) {
    TypeParam m;

    ASSERT_TRUE(m.try_lock_upgrade_for(std::chrono::milliseconds(10)));
    m.unlock_upgrade();
}

// 升级锁与读锁兼容
TYPED_TEST(UpgradeMutexTest, UpgradeLockCompatibleWithSharedLock) {
    TypeParam          m;
    std::promise<void> upgrade_locked_promise;
    auto               upgrade_locked_future = upgrade_locked_promise.get_future();
    std::promise<void> shared_locked_promise;
    auto               shared_locked_future = shared_locked_promise.get_future();
    std::atomic<bool>  both_locked{false};

    mc::test::runtime::future_flag release_flag;

    std::thread upgrade_thread([&] {
        m.lock_upgrade();
        upgrade_locked_promise.set_value();

        // 等待共享锁也获取成功
        shared_locked_future.wait();
        both_locked = true;

        // 等待释放信号，最多等待2秒
        if (!release_flag.wait_for(std::chrono::milliseconds(2000))) {
            m.unlock_upgrade();
            return; // 超时退出
        }
        m.unlock_upgrade();
    });

    std::thread shared_thread([&] {
        upgrade_locked_future.wait(); // 等待升级锁先获取

        m.lock_shared();
        shared_locked_promise.set_value();

        // 等待释放信号，最多等待2秒
        if (!release_flag.wait_for(std::chrono::milliseconds(2000))) {
            m.unlock_shared();
            return; // 超时退出
        }
        m.unlock_shared();
    });

    // 等待两个线程都获取锁
    std::this_thread::sleep_for(std::chrono::milliseconds(10));
    release_flag.set();
    upgrade_thread.join();
    shared_thread.join();

    ASSERT_TRUE(both_locked.load());
}

// 升级锁与写锁互斥
TYPED_TEST(UpgradeMutexTest, UpgradeLockMutuallyExclusiveWithWriteLock) {
    TypeParam          m;
    std::promise<void> upgrade_locked_promise;
    auto               upgrade_locked_future = upgrade_locked_promise.get_future();
    std::atomic<bool>  write_attempted{false};
    std::atomic<bool>  write_succeeded{false};

    mc::test::runtime::future_flag release_flag;

    std::thread upgrade_thread([&] {
        m.lock_upgrade();
        upgrade_locked_promise.set_value();
        // 等待释放信号，最多等待2秒
        if (!release_flag.wait_for(std::chrono::milliseconds(2000))) {
            m.unlock_upgrade();
            return; // 超时退出
        }
        m.unlock_upgrade();
    });

    std::thread write_thread([&] {
        upgrade_locked_future.wait(); // 等待升级锁先获取
        write_attempted = true;

        if (m.try_lock_for(std::chrono::milliseconds(10))) {
            write_succeeded = true;
            m.unlock();
        }
    });

    // 等待写线程完成
    write_thread.join();
    release_flag.set();
    upgrade_thread.join();

    ASSERT_TRUE(write_attempted.load());
    ASSERT_FALSE(write_succeeded.load()); // 写锁应该失败
}

// 测试从升级锁升级到写锁
TYPED_TEST(UpgradeMutexTest, UpgradeToWriteLock) {
    TypeParam m;

    m.lock_upgrade();
    m.unlock_upgrade_and_lock(); // 升级到写锁
    m.unlock();
}

// 测试从升级锁升级到写锁的超时功能
TYPED_TEST(UpgradeMutexTest, UpgradeToWriteLockWithTimeout) {
    TypeParam          m;
    std::promise<void> reader_ready_promise;
    auto               reader_ready_future = reader_ready_promise.get_future();
    std::atomic<bool>  upgrade_attempted{false};
    std::atomic<bool>  upgrade_succeeded{false};

    mc::test::runtime::future_flag release_flag;

    // 启动一个读线程持有读锁
    std::thread reader_thread([&] {
        m.lock_shared();
        reader_ready_promise.set_value();
        // 等待释放信号，最多等待2秒
        if (!release_flag.wait_for(std::chrono::milliseconds(2000))) {
            m.unlock_shared();
            return; // 超时退出
        }
        m.unlock_shared();
    });

    // 在另一个线程中尝试从升级锁升级到写锁
    bool upgrade_finished = false;
    std::thread upgrade_thread([&] {
        reader_ready_future.wait();

        m.lock_upgrade(); // 这应该成功，因为升级锁与读锁兼容
        upgrade_attempted = true;

        // 尝试升级到写锁，但由于读锁存在应该超时
        if (m.try_unlock_upgrade_and_lock_for(std::chrono::milliseconds(50))) {
            upgrade_succeeded = true;
            m.unlock();
        } else {
            m.unlock_upgrade(); // 超时后释放升级锁
        }
        upgrade_finished = true;
    });

    if (!upgrade_thread.joinable()) {
        ADD_FAILURE() << "Upgrade thread was not joinable";
    } else {
        // 等待升级线程完成，最多等待2秒
        constexpr auto wait_limit = std::chrono::milliseconds(2000);
        auto           start      = std::chrono::steady_clock::now();
        while (!upgrade_finished &&
               std::chrono::steady_clock::now() - start < wait_limit) {
            std::this_thread::yield();
        }
        if (!upgrade_finished) {
            release_flag.set(); // 确保读线程退出
            reader_thread.join();
            GTEST_SKIP() << "Upgrade thread did not finish in time under sanitizer";
        }
        release_flag.set(); // 释放读锁
        reader_thread.join();
        upgrade_thread.join();
    }

    ASSERT_TRUE(upgrade_attempted.load());
    ASSERT_FALSE(upgrade_succeeded.load()); // 升级应该超时失败
}

// 测试锁降级功能
TYPED_TEST(UpgradeMutexTest, DowngradeWriteToUpgrade) {
    TypeParam m;

    m.lock();
    m.unlock_and_lock_upgrade(); // 从写锁降级到升级锁
    m.unlock_upgrade();
}

TYPED_TEST(UpgradeMutexTest, DowngradeWriteToShared) {
    TypeParam m;

    m.lock();
    m.unlock_and_lock_shared(); // 从写锁降级到读锁
    m.unlock_shared();
}

TYPED_TEST(UpgradeMutexTest, DowngradeUpgradeToShared) {
    TypeParam m;

    m.lock_upgrade();
    m.unlock_upgrade_and_lock_shared(); // 从升级锁降级到读锁
    m.unlock_shared();
}

// 复杂的并发场景测试
TYPED_TEST(UpgradeMutexTest, ComplexConcurrentScenario) {
    TypeParam m;

    struct SharedData {
        long long value      = 0;
        long long validation = 0;
    };

    SharedData               data;
    std::atomic<bool>        running{true};
    constexpr int           num_readers          = 3;
    constexpr int           reader_iterations    = 50;
    constexpr int           writer_iterations    = 100;
    constexpr int           upgrader_iterations  = 50;
    std::vector<std::thread> reader_threads;
    std::vector<std::thread> writer_threads;
    std::vector<std::thread> upgrader_threads;
    std::vector<std::promise<void>> reader_done_promises(num_readers);
    std::vector<std::future<void>>  reader_done_futures;
    reader_done_futures.reserve(num_readers);
    for (auto& promise : reader_done_promises) {
        reader_done_futures.push_back(promise.get_future());
    }

    // 读线程
    for (int idx = 0; idx < num_readers; ++idx) {
        reader_threads.emplace_back([&, idx]() {
            for (int iter = 0; iter < reader_iterations; ++iter) {
                std::shared_lock<TypeParam> lock(m);
                ASSERT_EQ(data.value, data.validation);
                std::this_thread::yield();
            }
            reader_done_promises[idx].set_value();
        });
    }

    // 写线程
    writer_threads.emplace_back([&]() {
        for (int i = 0; i < writer_iterations; ++i) {
            std::lock_guard<TypeParam> lock(m);
            data.value += 2;
            data.validation = data.value;
        }
    });

    // 升级锁线程
    upgrader_threads.emplace_back([&]() {
        for (int i = 0; i < upgrader_iterations; ++i) {
            if (!m.try_lock_upgrade_for(std::chrono::milliseconds(20))) {
                continue;
            }

            long long current_value = data.value;
            bool      promoted      = false;
            if (i % 2 == 0 || current_value % 2 == 0) {
                if (m.try_unlock_upgrade_and_lock_for(std::chrono::milliseconds(20))) {
                    data.value += 3;
                    data.validation = data.value;
                    m.unlock();
                    promoted = true;
                }
            }
            if (!promoted) {
                m.unlock_upgrade();
            }
        }
    });

    auto join_with_timeout = [](std::thread& t, const char* name) {
        if (!t.joinable()) {
            return;
        }
        auto start = std::chrono::steady_clock::now();
        while (true) {
            using namespace std::chrono_literals;
            if (std::chrono::steady_clock::now() - start > 500ms) {
                t.detach();
                GTEST_SKIP() << name << " thread did not finish within timeout";
                return;
            }
            if (t.joinable()) {
                if (t.joinable()) {
                    t.join();
                    break;
                }
            }
        }
    };

    join_with_timeout(writer_threads.front(), "Writer");
    join_with_timeout(upgrader_threads.front(), "Upgrader");

    running = false;

    for (auto& future : reader_done_futures) {
        if (future.wait_for(std::chrono::milliseconds(200)) != std::future_status::ready) {
            GTEST_SKIP() << "Reader threads did not finish within timeout";
        }
    }
    for (auto& thread : reader_threads) {
        if (thread.joinable()) {
            thread.join();
        }
    }

    // 验证最终状态一致性
    ASSERT_EQ(data.value, data.validation);
    if constexpr (std::is_same_v<TypeParam, reader_priority_shared_mutex>) {
        ASSERT_GT(data.value, writer_iterations * 2); // 读优先策略应该包含升级锁线程的修改
    } else {
        ASSERT_GE(data.value, writer_iterations * 2); // 写优先策略至少包含写线程的修改
    }
}
